#include <filesystem>
#include <fairlogger/Logger.h>

void draw_checkAlignment(const char* fileName = "Unpack.root") {
    fair::Logger::SetConsoleColor(true);
    // fair::Logger::SetConsoleSeverity("debug"); 

    // ====== 打开 ROOT 文件 ======
    TFile* fFile = TFile::Open(fileName);
    if (!fFile || fFile->IsZombie()) {
        LOG(error) << "Cannot open file.";
        return;
    }

    // ====== 创建Tree及Data类 ======
    TTree* fTree = (TTree*) fFile->Get("HiTree");

    TClonesArray* fTofData{nullptr};
    TClonesArray* fTpcData{nullptr};
    TClonesArray* fMusicData{nullptr};

    fTree->SetBranchAddress("HiTofUnpackedData", &fTofData);
    fTree->SetBranchAddress("HiTpcUnpackedData", &fTpcData);
    fTree->SetBranchAddress("HiMusicUnpackedData", &fMusicData);

    // ====== 查找第一个有数据的时间片 ======
    int TSIndex = -1;
    int nTofEntries = 0, nTpcEntries = 0, nMusicEntries = 0;

    for(int i = 0; i < fTree->GetEntries(); ++i) {
        fTree->GetEntry(i);
        if (fTofData->GetEntriesFast() || fTpcData->GetEntriesFast() ||
             fMusicData->GetEntriesFast()) {
            TSIndex = i;
            break;
        }
    }

    LOG(info)<<"Draw since TS No."<<TSIndex;
    // ====== 画每个时间片 ======
    gStyle->SetOptStat(0);
    gStyle->SetLabelFont(62, "XYZ");
    gStyle->SetTitleFont(62, "XYZ");
    gStyle->SetTitleFont(62, "");
    gStyle->SetStatFont(62);
    gStyle->SetLabelSize(0.05, "XYZ");
    gStyle->SetTitleSize(0.05, "XYZ");
    // gStyle->SetPalette(kRainBow); 
    gStyle->SetPalette(1); 

    TCanvas* c = new TCanvas("c", "Detector Alignment", 1600, 900);
    // c->Divide(1, 3, 0, 0); // 垂直三块（无间距）
    c->Divide(1, 3); // 垂直三块（无间距）
    
	const float sliceTimeUs = 102.4; // 每片持续时间
	int NTS = 20; // 画多少个时间片
	float totalTime = NTS * sliceTimeUs;
	//初始化波形数据
    std::vector<double> x_TPC(4096*NTS), y_TPC(4096*NTS, 0), x_Music(4096*NTS);
	std::vector<double> y_Music;

	for(int i = 0; i < 4096 * NTS; ++i){
		x_TPC[i] = x_Music[i] = i * 25. / 1000.; // 单位 us
	}

    TH1F* hTof = new TH1F("hTof", "TOF TimeOffset", 2000, 0, totalTime);
    
    for (int i = TSIndex; (i < TSIndex + NTS) && (i < fTree->GetEntries()); ++i) {
        float baseTime = (i - TSIndex) * sliceTimeUs;
        fTree->GetEntry(i);
        // ---- TOF ----
		if(fTofData->GetEntriesFast()){
            auto* tofdata = (HiTofUnpackedData*)fTofData->At(0);
			hTof->Fill(baseTime + tofdata->GetTimeOffset() * 0.025); 
            LOG(debug)<<"tofdata->GetTimeOffset() = "<<tofdata->GetTimeOffset();
		}
		// ---- TPC ----
		if(fTpcData->GetEntriesFast()){
            auto* tpcdata = (HiTpcUnpackedData*)fTpcData->At(0);
            for (size_t j = 0; j < tpcdata->GetWaveformData().size(); j++) {
                y_TPC[(i - TSIndex)*4096 + tpcdata->GetTimeOffset() + j] 
                    = tpcdata->GetWaveformData()[j];
			}
            LOG(debug)<<"tpcdata->GetTimeOffset() = "<<tpcdata->GetTimeOffset();
		}
		// ---- MUSIC ----
		if(fMusicData->GetEntriesFast()){
            auto* musicdata = (HiMusicUnpackedData*)fMusicData->At(0);
			for(auto& wf : musicdata->GetWaveformData()[5]){
				y_Music.emplace_back(wf);
			}
            LOG(debug)<<"draw music data";

		}
		else{
			for (size_t j = 0; j < 4096; j++) {
				y_Music.push_back(0);
			}
		}
    }
	c->cd(1);
	// gPad->SetBottomMargin(0);
	hTof->GetYaxis()->SetTitle("Counts");
	hTof->GetYaxis()->SetTitleOffset(0.5);
	hTof->GetYaxis()->CenterTitle();
	hTof->SetLineColor(kRed);
	hTof->SetLineWidth(2);
	hTof->Draw();
	TGraph* gMusic = new TGraph(x_Music.size(), x_Music.data(), y_Music.data());
	c->cd(2);
	// gPad->SetTopMargin(0);
    // gPad->SetBottomMargin(0);
	gMusic->SetTitle("MUSIC Waveform;;Amplitude");
	gMusic->GetYaxis()->SetTitleOffset(0.5);
	gMusic->GetYaxis()->CenterTitle();
	gMusic->SetLineColor(kBlue);
	gMusic->SetLineWidth(2);
	gMusic->GetXaxis()->SetRangeUser(0, totalTime);
	gMusic->GetYaxis()->SetRangeUser(8000, 8600);
	gMusic->Draw("AL");
	TGraph* gTpc = new TGraph(x_TPC.size(), x_TPC.data(), y_TPC.data());
	c->cd(3);
	// gPad->SetTopMargin(0);
	gTpc->SetTitle("TPC Waveform;Time [us];Amplitude");
	gTpc->GetYaxis()->SetTitleOffset(0.5);
	gTpc->GetYaxis()->CenterTitle();
	gTpc->SetLineColor(kGreen);
	gTpc->SetLineWidth(2);
	gTpc->GetXaxis()->SetRangeUser(0, totalTime);
	// gTpc->GetYaxis()->SetRangeUser(0, 10);
	gTpc->Draw("AL");
	// === 画时间片分界线 ===
    for (int i = 1; i < NTS; ++i) {
        double x = i * sliceTimeUs;

        c->cd(1);
        TLine* l1 = new TLine(x, 0, x, 5);
        l1->SetLineStyle(2);
        l1->Draw();

        c->cd(2);
        TLine* l2 = new TLine(x, 8000, x, 1e5);
        l2->SetLineStyle(2);
        l2->Draw();

        c->cd(3);
        TLine* l3 = new TLine(x, 0, x, 20);
        l3->SetLineStyle(2);
        l3->Draw();
    }
    // ====== 图例 ======
    // auto legend = new TLegend(0.8, 0.8, 0.95, 0.95);
    // legend->AddEntry(hTof, "TOF", "l");
    // legend->AddEntry(gMusic, "MUSIC", "lp");
    // legend->AddEntry(gTpc, "TPC", "lp");
    // legend->Draw();

    c->Update();
}

